import os

import cv2
from PIL import Image
from ultralyticsplus import YOLO, render_result

import baseconf
from baseconf import BASE_DISK

model_path="/model_path/yolov8s/yolov8s.pt" if os.path.exists(
                        "/etc") else BASE_DISK + ":/model_path/yolov8s/yolov8s.pt"
# Load a model
model = YOLO(model_path)  # load an official model



# load model


# set model parameters
# model.overrides['conf'] = 0.25  # NMS confidence threshold
# model.overrides['iou'] = 0.45  # NMS IoU threshold
# model.overrides['agnostic_nms'] = False  # NMS class-agnostic
# model.overrides['max_det'] = 1000  # maximum number of detections per image


fristvideo_path = baseconf.BASE_DISK+":/datasets_path/tianchivit/fristvideo"

cap = cv2.VideoCapture(fristvideo_path+'./43.avi')
for _ in range(30):
    cap.read()
img = cap.read()[1]
image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image = Image.fromarray(image)
# set image

print(model(image))

# perform inference
results = model.predict(image)

# observe results
print(results[0].boxes)
render = render_result(model=model, image=image, result=results[0])
render.show()
